#include "network_protocol.h"
#include "../memory/memory_manager.h"
#include "../auth/auth_manager.h"
#include "../parser/sql_parser.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <time.h>
#include <sys/time.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <signal.h>

// Global server capabilities
static const uint32_t SERVER_CAPABILITIES = 
    CLIENT_LONG_PASSWORD | CLIENT_FOUND_ROWS | CLIENT_LONG_FLAG |
    CLIENT_CONNECT_WITH_DB | CLIENT_ODBC | CLIENT_LOCAL_FILES |
    CLIENT_IGNORE_SPACE | CLIENT_PROTOCOL_41 | CLIENT_INTERACTIVE |
    CLIENT_TRANSACTIONS | CLIENT_SECURE_CONNECTION | CLIENT_MULTI_STATEMENTS |
    CLIENT_MULTI_RESULTS | CLIENT_PS_MULTI_RESULTS | CLIENT_PLUGIN_AUTH |
    CLIENT_CONNECT_ATTRS | CLIENT_SESSION_TRACK | CLIENT_DEPRECATE_EOF;

static const uint16_t SERVER_STATUS_DEFAULT = 
    SERVER_STATUS_AUTOCOMMIT;

// Packet Management Functions

mysql_packet_t* mysql_packet_create(uint32_t length, uint8_t sequence_id) {
    mysql_packet_t* packet = malloc(sizeof(mysql_packet_t));
    if (!packet) return NULL;
    
    packet->length = length;
    packet->sequence_id = sequence_id;
    packet->payload_size = length;
    
    if (length > 0) {
        packet->payload = malloc(length);
        if (!packet->payload) {
            free(packet);
            return NULL;
        }
        memset(packet->payload, 0, length);
    } else {
        packet->payload = NULL;
    }
    
    return packet;
}

void mysql_packet_destroy(mysql_packet_t* packet) {
    if (!packet) return;
    
    if (packet->payload) {
        free(packet->payload);
    }
    free(packet);
}

int mysql_packet_write_uint8(mysql_packet_t* packet, size_t offset, uint8_t value) {
    if (!packet || offset >= packet->payload_size) return -1;
    
    packet->payload[offset] = value;
    return 0;
}

int mysql_packet_write_uint16(mysql_packet_t* packet, size_t offset, uint16_t value) {
    if (!packet || offset + 1 >= packet->payload_size) return -1;
    
    packet->payload[offset] = value & 0xff;
    packet->payload[offset + 1] = (value >> 8) & 0xff;
    return 0;
}

int mysql_packet_write_uint24(mysql_packet_t* packet, size_t offset, uint32_t value) {
    if (!packet || offset + 2 >= packet->payload_size) return -1;
    
    packet->payload[offset] = value & 0xff;
    packet->payload[offset + 1] = (value >> 8) & 0xff;
    packet->payload[offset + 2] = (value >> 16) & 0xff;
    return 0;
}

int mysql_packet_write_uint32(mysql_packet_t* packet, size_t offset, uint32_t value) {
    if (!packet || offset + 3 >= packet->payload_size) return -1;
    
    packet->payload[offset] = value & 0xff;
    packet->payload[offset + 1] = (value >> 8) & 0xff;
    packet->payload[offset + 2] = (value >> 16) & 0xff;
    packet->payload[offset + 3] = (value >> 24) & 0xff;
    return 0;
}

int mysql_packet_write_uint64(mysql_packet_t* packet, size_t offset, uint64_t value) {
    if (!packet || offset + 7 >= packet->payload_size) return -1;
    
    for (int i = 0; i < 8; i++) {
        packet->payload[offset + i] = (value >> (i * 8)) & 0xff;
    }
    return 0;
}

int mysql_packet_write_string(mysql_packet_t* packet, size_t offset, const char* str) {
    if (!packet || !str) return -1;
    
    size_t len = strlen(str);
    if (offset + len > packet->payload_size) return -1;
    
    memcpy(packet->payload + offset, str, len);
    return len;
}

int mysql_packet_write_lenenc_int(mysql_packet_t* packet, size_t offset, uint64_t value) {
    if (!packet) return -1;
    
    if (value < 0xfb) {
        if (offset >= packet->payload_size) return -1;
        packet->payload[offset] = (uint8_t)value;
        return 1;
    } else if (value < 0x10000) {
        if (offset + 2 >= packet->payload_size) return -1;
        packet->payload[offset] = 0xfc;
        mysql_packet_write_uint16(packet, offset + 1, (uint16_t)value);
        return 3;
    } else if (value < 0x1000000) {
        if (offset + 3 >= packet->payload_size) return -1;
        packet->payload[offset] = 0xfd;
        mysql_packet_write_uint24(packet, offset + 1, (uint32_t)value);
        return 4;
    } else {
        if (offset + 8 >= packet->payload_size) return -1;
        packet->payload[offset] = 0xfe;
        mysql_packet_write_uint64(packet, offset + 1, value);
        return 9;
    }
}

int mysql_packet_write_lenenc_string(mysql_packet_t* packet, size_t offset, const char* str) {
    if (!packet || !str) return -1;
    
    size_t str_len = strlen(str);
    int len_bytes = mysql_packet_write_lenenc_int(packet, offset, str_len);
    if (len_bytes < 0) return -1;
    
    if (offset + len_bytes + str_len > packet->payload_size) return -1;
    
    memcpy(packet->payload + offset + len_bytes, str, str_len);
    return len_bytes + str_len;
}

// Read functions
uint8_t mysql_packet_read_uint8(mysql_packet_t* packet, size_t offset) {
    if (!packet || offset >= packet->payload_size) return 0;
    return packet->payload[offset];
}

uint16_t mysql_packet_read_uint16(mysql_packet_t* packet, size_t offset) {
    if (!packet || offset + 1 >= packet->payload_size) return 0;
    return packet->payload[offset] | (packet->payload[offset + 1] << 8);
}

uint32_t mysql_packet_read_uint24(mysql_packet_t* packet, size_t offset) {
    if (!packet || offset + 2 >= packet->payload_size) return 0;
    return packet->payload[offset] | (packet->payload[offset + 1] << 8) | 
           (packet->payload[offset + 2] << 16);
}

uint32_t mysql_packet_read_uint32(mysql_packet_t* packet, size_t offset) {
    if (!packet || offset + 3 >= packet->payload_size) return 0;
    return packet->payload[offset] | (packet->payload[offset + 1] << 8) |
           (packet->payload[offset + 2] << 16) | (packet->payload[offset + 3] << 24);
}

uint64_t mysql_packet_read_uint64(mysql_packet_t* packet, size_t offset) {
    if (!packet || offset + 7 >= packet->payload_size) return 0;
    
    uint64_t result = 0;
    for (int i = 0; i < 8; i++) {
        result |= ((uint64_t)packet->payload[offset + i]) << (i * 8);
    }
    return result;
}

uint64_t mysql_packet_read_lenenc_int(mysql_packet_t* packet, size_t offset) {
    if (!packet || offset >= packet->payload_size) return 0;
    
    uint8_t first_byte = packet->payload[offset];
    
    if (first_byte < 0xfb) {
        return first_byte;
    } else if (first_byte == 0xfc) {
        return mysql_packet_read_uint16(packet, offset + 1);
    } else if (first_byte == 0xfd) {
        return mysql_packet_read_uint24(packet, offset + 1);
    } else if (first_byte == 0xfe) {
        return mysql_packet_read_uint64(packet, offset + 1);
    }
    
    return 0;
}

char* mysql_packet_read_string(mysql_packet_t* packet, size_t offset, size_t length) {
    if (!packet || offset + length > packet->payload_size) return NULL;
    
    char* str = malloc(length + 1);
    if (!str) return NULL;
    
    memcpy(str, packet->payload + offset, length);
    str[length] = '\0';
    return str;
}

char* mysql_packet_read_lenenc_string(mysql_packet_t* packet, size_t offset, uint64_t* length) {
    if (!packet || !length) return NULL;
    
    uint64_t str_len = mysql_packet_read_lenenc_int(packet, offset);
    size_t len_bytes = mysql_lenenc_int_size(str_len);
    
    *length = str_len;
    return mysql_packet_read_string(packet, offset + len_bytes, str_len);
}

// Network I/O Functions

int mysql_send_packet(mysql_connection_t* conn, mysql_packet_t* packet) {
    if (!conn || !packet) return -1;
    
    // Create wire format: [3 bytes length][1 byte sequence_id][payload]
    uint8_t header[4];
    header[0] = packet->length & 0xff;
    header[1] = (packet->length >> 8) & 0xff;
    header[2] = (packet->length >> 16) & 0xff;
    header[3] = packet->sequence_id;
    
    // Send header
    ssize_t sent = send(conn->sockfd, header, 4, 0);
    if (sent != 4) {
        return -1;
    }
    
    // Send payload if present
    if (packet->length > 0 && packet->payload) {
        sent = send(conn->sockfd, packet->payload, packet->length, 0);
        if (sent != packet->length) {
            return -1;
        }
    }
    
    conn->bytes_sent += 4 + packet->length;
    conn->sequence_id = (packet->sequence_id + 1) & 0xff;
    
    return 0;
}

mysql_packet_t* mysql_recv_packet(mysql_connection_t* conn) {
    if (!conn) return NULL;
    
    // Read packet header
    uint8_t header[4];
    ssize_t received = recv(conn->sockfd, header, 4, MSG_WAITALL);
    if (received != 4) {
        return NULL;
    }
    
    // Parse packet length and sequence ID
    uint32_t length = header[0] | (header[1] << 8) | (header[2] << 16);
    uint8_t sequence_id = header[3];
    
    // Validate packet length
    if (length > MYSQL_MAX_PACKET_SIZE) {
        return NULL;
    }
    
    // Create packet
    mysql_packet_t* packet = mysql_packet_create(length, sequence_id);
    if (!packet) return NULL;
    
    // Read payload if present
    if (length > 0) {
        received = recv(conn->sockfd, packet->payload, length, MSG_WAITALL);
        if (received != length) {
            mysql_packet_destroy(packet);
            return NULL;
        }
    }
    
    conn->bytes_received += 4 + length;
    return packet;
}

int mysql_send_handshake(mysql_connection_t* conn) {
    if (!conn) return -1;
    
    // Generate auth plugin data (salt)
    mysql_generate_auth_data(conn->auth_plugin_data);
    
    // Calculate packet size
    size_t server_version_len = strlen(MYSQL_SERVER_VERSION);
    size_t packet_size = 1 + server_version_len + 1 + 4 + 8 + 1 + 2 + 1 + 2 + 2 + 1 + 10 + 12 + 1 + 21 + 1;
    
    mysql_packet_t* packet = mysql_packet_create(packet_size, 0);
    if (!packet) return -1;
    
    size_t offset = 0;
    
    // Protocol version
    mysql_packet_write_uint8(packet, offset++, MYSQL_PROTOCOL_VERSION);
    
    // Server version
    memcpy(packet->payload + offset, MYSQL_SERVER_VERSION, server_version_len);
    offset += server_version_len;
    packet->payload[offset++] = 0; // null terminator
    
    // Connection ID
    mysql_packet_write_uint32(packet, offset, conn->connection_id);
    offset += 4;
    
    // Auth plugin data part 1 (8 bytes)
    memcpy(packet->payload + offset, conn->auth_plugin_data, 8);
    offset += 8;
    
    // Filler
    packet->payload[offset++] = 0;
    
    // Capability flags (lower 2 bytes)
    mysql_packet_write_uint16(packet, offset, SERVER_CAPABILITIES & 0xffff);
    offset += 2;
    
    // Character set
    packet->payload[offset++] = 0x21; // utf8_general_ci
    
    // Status flags
    mysql_packet_write_uint16(packet, offset, SERVER_STATUS_DEFAULT);
    offset += 2;
    
    // Capability flags (upper 2 bytes)
    mysql_packet_write_uint16(packet, offset, (SERVER_CAPABILITIES >> 16) & 0xffff);
    offset += 2;
    
    // Auth plugin data length
    packet->payload[offset++] = MYSQL_AUTH_PLUGIN_DATA_LEN;
    
    // Reserved (10 bytes)
    memset(packet->payload + offset, 0, 10);
    offset += 10;
    
    // Auth plugin data part 2 (12 bytes)
    memcpy(packet->payload + offset, conn->auth_plugin_data + 8, 12);
    offset += 12;
    
    // Null terminator for auth plugin data part 2
    packet->payload[offset++] = 0;
    
    // Auth plugin name
    const char* plugin_name = "mysql_native_password";
    size_t plugin_name_len = strlen(plugin_name);
    memcpy(packet->payload + offset, plugin_name, plugin_name_len);
    offset += plugin_name_len;
    packet->payload[offset++] = 0; // null terminator
    
    int result = mysql_send_packet(conn, packet);
    mysql_packet_destroy(packet);
    
    if (result == 0) {
        conn->state = CONN_STATE_HANDSHAKE_SENT;
    }
    
    return result;
}

int mysql_recv_handshake_response(mysql_connection_t* conn, mysql_handshake_response_t* response) {
    if (!conn || !response) return -1;
    
    mysql_packet_t* packet = mysql_recv_packet(conn);
    if (!packet) return -1;
    
    size_t offset = 0;
    
    // Capability flags
    response->capability_flags = mysql_packet_read_uint32(packet, offset);
    offset += 4;
    
    // Max packet size
    response->max_packet_size = mysql_packet_read_uint32(packet, offset);
    offset += 4;
    
    // Character set
    response->character_set = mysql_packet_read_uint8(packet, offset);
    offset += 1;
    
    // Reserved (23 bytes)
    offset += 23;
    
    // Username (null-terminated)
    size_t username_len = 0;
    while (offset + username_len < packet->length && 
           packet->payload[offset + username_len] != 0) {
        username_len++;
    }
    
    if (username_len >= sizeof(response->username)) {
        mysql_packet_destroy(packet);
        return -1;
    }
    
    memcpy(response->username, packet->payload + offset, username_len);
    response->username[username_len] = '\0';
    offset += username_len + 1; // +1 for null terminator
    
    // Auth response length and data
    if (response->capability_flags & CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA) {
        uint64_t auth_len = mysql_packet_read_lenenc_int(packet, offset);
        offset += mysql_lenenc_int_size(auth_len);
        response->auth_response_len = auth_len;
        
        if (auth_len > sizeof(response->auth_response)) {
            mysql_packet_destroy(packet);
            return -1;
        }
        
        memcpy(response->auth_response, packet->payload + offset, auth_len);
        offset += auth_len;
    } else if (response->capability_flags & CLIENT_SECURE_CONNECTION) {
        response->auth_response_len = mysql_packet_read_uint8(packet, offset);
        offset += 1;
        
        // Note: auth_response_len is uint8_t (max 255) and auth_response is 256 bytes,
        // so no bounds check needed here
        memcpy(response->auth_response, packet->payload + offset, response->auth_response_len);
        offset += response->auth_response_len;
    }
    
    // Database (optional)
    if ((response->capability_flags & CLIENT_CONNECT_WITH_DB) && offset < packet->length) {
        size_t db_len = 0;
        while (offset + db_len < packet->length && 
               packet->payload[offset + db_len] != 0) {
            db_len++;
        }
        
        if (db_len < sizeof(response->database)) {
            memcpy(response->database, packet->payload + offset, db_len);
            response->database[db_len] = '\0';
            offset += db_len + 1;
        }
    }
    
    // Auth plugin name (optional)
    if ((response->capability_flags & CLIENT_PLUGIN_AUTH) && offset < packet->length) {
        size_t plugin_len = 0;
        while (offset + plugin_len < packet->length && 
               packet->payload[offset + plugin_len] != 0) {
            plugin_len++;
        }
        
        if (plugin_len < sizeof(response->auth_plugin_name)) {
            memcpy(response->auth_plugin_name, packet->payload + offset, plugin_len);
            response->auth_plugin_name[plugin_len] = '\0';
        }
    }
    
    mysql_packet_destroy(packet);
    return 0;
}

int mysql_send_ok(mysql_connection_t* conn, uint64_t affected_rows, uint64_t last_insert_id) {
    if (!conn) return -1;
    
    // Calculate packet size
    size_t packet_size = 1 + mysql_lenenc_int_size(affected_rows) + 
                        mysql_lenenc_int_size(last_insert_id) + 2 + 2;
    
    if (conn->capability_flags & CLIENT_SESSION_TRACK) {
        packet_size += mysql_lenenc_int_size(0); // session state changes
    }
    
    mysql_packet_t* packet = mysql_packet_create(packet_size, conn->sequence_id);
    if (!packet) return -1;
    
    size_t offset = 0;
    
    // OK header
    mysql_packet_write_uint8(packet, offset++, MYSQL_TYPE_OK);
    
    // Affected rows
    offset += mysql_packet_write_lenenc_int(packet, offset, affected_rows);
    
    // Last insert ID
    offset += mysql_packet_write_lenenc_int(packet, offset, last_insert_id);
    
    // Status flags
    mysql_packet_write_uint16(packet, offset, conn->status_flags);
    offset += 2;
    
    // Warnings
    mysql_packet_write_uint16(packet, offset, 0);
    offset += 2;
    
    // Session state changes (if supported)
    if (conn->capability_flags & CLIENT_SESSION_TRACK) {
        mysql_packet_write_lenenc_int(packet, offset, 0);
    }
    
    int result = mysql_send_packet(conn, packet);
    mysql_packet_destroy(packet);
    
    return result;
}

int mysql_send_error(mysql_connection_t* conn, uint16_t error_code, 
                    const char* sql_state, const char* message) {
    if (!conn || !message) return -1;
    
    size_t message_len = strlen(message);
    size_t packet_size = 1 + 2 + 1 + 5 + message_len;
    
    mysql_packet_t* packet = mysql_packet_create(packet_size, conn->sequence_id);
    if (!packet) return -1;
    
    size_t offset = 0;
    
    // Error header
    mysql_packet_write_uint8(packet, offset++, MYSQL_TYPE_ERR);
    
    // Error code
    mysql_packet_write_uint16(packet, offset, error_code);
    offset += 2;
    
    // SQL state marker
    packet->payload[offset++] = '#';
    
    // SQL state
    const char* state = sql_state ? sql_state : "HY000";
    memcpy(packet->payload + offset, state, 5);
    offset += 5;
    
    // Error message
    memcpy(packet->payload + offset, message, message_len);
    
    int result = mysql_send_packet(conn, packet);
    mysql_packet_destroy(packet);
    
    return result;
}

int mysql_send_eof(mysql_connection_t* conn, uint16_t warnings, uint16_t status_flags) {
    if (!conn) return -1;
    
    mysql_packet_t* packet = mysql_packet_create(5, conn->sequence_id);
    if (!packet) return -1;
    
    mysql_packet_write_uint8(packet, 0, MYSQL_TYPE_EOF);
    mysql_packet_write_uint16(packet, 1, warnings);
    mysql_packet_write_uint16(packet, 3, status_flags);
    
    int result = mysql_send_packet(conn, packet);
    mysql_packet_destroy(packet);
    
    return result;
}

// Authentication Functions

void mysql_generate_auth_data(char* auth_data) {
    if (!auth_data) return;
    
    srand(time(NULL));
    for (int i = 0; i < MYSQL_AUTH_PLUGIN_DATA_LEN; i++) {
        auth_data[i] = (char)(rand() % 256);
        // Ensure printable characters for compatibility
        if (auth_data[i] == 0 || (unsigned char)auth_data[i] == 255) {
            auth_data[i] = (char)((rand() % 94) + 33); // ASCII 33-126
        }
    }
    auth_data[MYSQL_AUTH_PLUGIN_DATA_LEN] = '\0';
}

int mysql_authenticate_user(mysql_connection_t* conn, const char* username, 
                           const char* auth_response, const char* database) {
    if (!conn || !username) return -1;
    
    // For now, implement basic authentication
    // In a full implementation, this would check against auth_manager
    // Note: auth_response would be used for password verification
    (void)auth_response;  // Suppress unused parameter warning
    
    // Store authentication info
    strncpy(conn->username, username, sizeof(conn->username) - 1);
    conn->username[sizeof(conn->username) - 1] = '\0';
    
    if (database) {
        strncpy(conn->database, database, sizeof(conn->database) - 1);
        conn->database[sizeof(conn->database) - 1] = '\0';
    }
    
    conn->authenticated = true;
    conn->state = CONN_STATE_AUTHENTICATED;
    
    return 0;
}

// Connection Management

mysql_connection_t* mysql_connection_create(int sockfd, struct sockaddr_in* client_addr) {
    mysql_connection_t* conn = malloc(sizeof(mysql_connection_t));
    if (!conn) return NULL;
    
    memset(conn, 0, sizeof(mysql_connection_t));
    
    conn->sockfd = sockfd;
    if (client_addr) {
        conn->client_addr = *client_addr;
    }
    
    conn->state = CONN_STATE_INIT;
    conn->sequence_id = 0;
    conn->capability_flags = 0;
    conn->character_set = 0x21; // utf8_general_ci
    conn->status_flags = SERVER_STATUS_DEFAULT;
    conn->authenticated = false;
    conn->in_transaction = false;
    conn->autocommit = true;
    
    // Initialize statistics
    conn->connect_time = time(NULL);
    conn->last_activity = conn->connect_time;
    
    // Initialize mutex
    if (pthread_mutex_init(&conn->mutex, NULL) != 0) {
        free(conn);
        return NULL;
    }
    
    return conn;
}

void mysql_connection_destroy(mysql_connection_t* conn) {
    if (!conn) return;
    
    // Close socket
    if (conn->sockfd >= 0) {
        close(conn->sockfd);
    }
    
    // Clean up transaction
    if (conn->current_txn) {
        // In a full implementation, roll back transaction
        conn->current_txn = NULL;
    }
    
    // Clean up buffers
    mysql_connection_cleanup_buffers(conn);
    
    // Destroy mutex
    pthread_mutex_destroy(&conn->mutex);
    
    free(conn);
}

int mysql_connection_init_buffers(mysql_connection_t* conn, size_t read_size, size_t write_size) {
    if (!conn) return -1;
    
    conn->read_buffer = malloc(read_size);
    if (!conn->read_buffer) return -1;
    
    conn->write_buffer = malloc(write_size);
    if (!conn->write_buffer) {
        free(conn->read_buffer);
        conn->read_buffer = NULL;
        return -1;
    }
    
    conn->read_buffer_size = read_size;
    conn->write_buffer_size = write_size;
    conn->read_buffer_pos = 0;
    conn->write_buffer_pos = 0;
    
    return 0;
}

void mysql_connection_cleanup_buffers(mysql_connection_t* conn) {
    if (!conn) return;
    
    if (conn->read_buffer) {
        free(conn->read_buffer);
        conn->read_buffer = NULL;
    }
    
    if (conn->write_buffer) {
        free(conn->write_buffer);
        conn->write_buffer = NULL;
    }
    
    conn->read_buffer_size = 0;
    conn->write_buffer_size = 0;
    conn->read_buffer_pos = 0;
    conn->write_buffer_pos = 0;
}

// Utility Functions

size_t mysql_lenenc_int_size(uint64_t value) {
    if (value < 0xfb) return 1;
    else if (value < 0x10000) return 3;
    else if (value < 0x1000000) return 4;
    else return 9;
}

size_t mysql_lenenc_string_size(const char* str) {
    if (!str) return mysql_lenenc_int_size(0);
    
    size_t len = strlen(str);
    return mysql_lenenc_int_size(len) + len;
}

const char* mysql_get_command_name(mysql_command_t command) {
    switch (command) {
        case COM_QUIT: return "COM_QUIT";
        case COM_INIT_DB: return "COM_INIT_DB";
        case COM_QUERY: return "COM_QUERY";
        case COM_PING: return "COM_PING";
        case COM_STATISTICS: return "COM_STATISTICS";
        case COM_PROCESS_INFO: return "COM_PROCESS_INFO";
        case COM_STMT_PREPARE: return "COM_STMT_PREPARE";
        case COM_STMT_EXECUTE: return "COM_STMT_EXECUTE";
        case COM_STMT_CLOSE: return "COM_STMT_CLOSE";
        default: return "UNKNOWN_COMMAND";
    }
}

const char* mysql_get_error_message(uint16_t error_code) {
    switch (error_code) {
        case ER_ACCESS_DENIED_ERROR: return "Access denied";
        case ER_NO_DB_ERROR: return "No database selected";
        case ER_UNKNOWN_COM_ERROR: return "Unknown command";
        case ER_SYNTAX_ERROR: return "Syntax error";
        case ER_TABLE_EXISTS_ERROR: return "Table already exists";
        case ER_BAD_TABLE_ERROR: return "Table doesn't exist";
        default: return "Unknown error";
    }
}